import pdfplumber
from PyPDF2 import PdfReader, PdfWriter

def extract_titles(pdf_path, font_size_threshold=14, bold_weight_threshold=700):
    """
    使用 pdfplumber 提取 PDF 中的标题，并支持多级目录。
    
    :param pdf_path: 输入 PDF 文件路径
    :param font_size_threshold: 字体大小阈值，大于此值的文本可能为标题
    :param bold_weight_threshold: 加粗权重阈值，大于此值的文本可能为加粗
    :return: 标题列表 [(页码, 标题, 父级索引), ...]
    """
    titles = []

    with pdfplumber.open(pdf_path) as pdf:
        for page_num, page in enumerate(pdf.pages):
            # 提取当前页的所有字符
            chars = page.chars
            if not chars:
                continue

            # 遍历字符，按行分组
            lines = {}
            for char in chars:
                top = round(char["top"], 1)  # 按行分组
                if top not in lines:
                    lines[top] = []
                lines[top].append(char)

            # 分析每一行
            for top, line_chars in lines.items():
                text = "".join([char["text"] for char in line_chars]).strip()
                if not text:  # 跳过空标题
                    continue

                font_sizes = [char["size"] for char in line_chars]
                avg_font_size = sum(font_sizes) / len(font_sizes)
                is_bold = any(char.get("fontname", "").endswith("Bold") for char in line_chars)

                # 判断是否为标题
                if avg_font_size > font_size_threshold or is_bold:
                    # 根据字体大小判断层级
                    level = 0 if avg_font_size >= font_size_threshold + 4 else 1
                    titles.append((page_num, text, level))

    return titles


def organize_titles(titles):
    """
    组织标题为多级目录结构。
    
    :param titles: 标题列表 [(页码, 标题, 层级), ...]
    :return: 多级目录 [(页码, 标题, 父级索引), ...]
    """
    organized_titles = []
    parent_indices = {}  # 记录每个层级的父级索引

    for page_num, title, level in titles:
        # 找到当前层级的父级索引
        parent_index = parent_indices.get(level - 1, None)
        organized_titles.append((page_num, title, parent_index))
        # 更新当前层级的父级索引
        parent_indices[level] = len(organized_titles) - 1

    return organized_titles


def add_bookmarks(input_pdf_path, output_pdf_path, bookmarks):
    """
    为 PDF 文件添加书签并保存。
    
    :param input_pdf_path: 输入 PDF 文件路径
    :param output_pdf_path: 输出 PDF 文件路径
    :param bookmarks: 书签列表 [(页码, 标题, 父级索引), ...]
    """
    reader = PdfReader(input_pdf_path)
    writer = PdfWriter()

    # 添加所有页面到 writer
    for page in reader.pages:
        writer.add_page(page)

    # 添加书签
    parent_bookmarks = {}
    for page_num, title, parent_index in bookmarks:
        if parent_index is None:
            # 添加顶级书签
            bookmark = writer.add_outline_item(title, page_num)
            parent_bookmarks[title] = bookmark
        else:
            # 添加子书签
            parent_title = bookmarks[parent_index][1]
            parent = parent_bookmarks[parent_title]
            writer.add_outline_item(title, page_num, parent=parent)

    # 保存带有书签的 PDF 文件
    with open(output_pdf_path, "wb") as output_pdf_file:
        writer.write(output_pdf_file)


# 示例用法
if __name__ == "__main__":
    # 输入文件路径
    input_pdf = "/Users/lingyu/study/book/computer/python/鸢尾花/Book1_Python-For-Beginners/Book_1_Ch_07__Python控制结构__编程不难__鸢尾花书__从加减乘除到机器学习.pdf"
    # 输出文件路径
    output_pdf = "example_with_auto_generated_bookmarks.pdf"

    # 提取标题
    raw_titles = extract_titles(input_pdf)

    # 组织为多级目录
    organized_titles = organize_titles(raw_titles)

    # 打印提取的标题
    print("提取的标题：")
    for page_num, title, parent_index in organized_titles:
        print(f"Page {page_num + 1}: {title} (Parent Index: {parent_index})")

    # 添加书签
    add_bookmarks(input_pdf, output_pdf, organized_titles)

    print(f"已生成带书签的 PDF 文件：{output_pdf}")